package com.demo.spark.rdd;

import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.api.java.function.Function2;

import java.util.Arrays;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;

public class RddMapPartitionWithIndex {
    public static void main(String[] args) {

        //RDD有两个分区
        SparkConf config = new SparkConf().setAppName("rdd map partition with index").setMaster("local");
        JavaSparkContext javaSparkContext = new JavaSparkContext(config);
        javaSparkContext.setLogLevel("error");
        List<Integer> data = Arrays.asList(1, 2, 4, 3, 5, 6, 7);
        JavaRDD<Integer> javaRDD = javaSparkContext.parallelize(data,2);
        //分区index、元素值、元素编号输出

        JavaRDD<String> mapPartitionsWithIndexRDD =
                javaRDD.mapPartitionsWithIndex(new Function2<Integer, Iterator<Integer>, Iterator<String>>() {
            @Override
            public Iterator<String> call(Integer v1, Iterator<Integer> v2) throws Exception {
                LinkedList<String> linkedList = new LinkedList<String>();
                int i = 0;
                while (v2.hasNext())
                    linkedList.add("partition:"+Integer.toString(v1) + "|" + v2.next().toString() + ", rownum:" + Integer.toString(i++));
                return linkedList.iterator();
            }
        },false);

        for (String s : mapPartitionsWithIndexRDD.collect()) {
            System.out.println(s);
        }
    }
}
